# 失真工具函数
import numpy as np
import cv2
from PIL import Image
import bm3d
from scipy import ndimage

def add_gaussian_noise(img, sigma):
    """添加高斯噪声"""
    noise = np.random.normal(0, sigma * 255, img.shape).astype(np.float32)
    noisy_img = np.clip(img.astype(np.float32) + noise, 0, 255).astype(np.uint8)
    return noisy_img

def add_gaussian_noise_color(img, sigma):
    """在YCbCr颜色空间中添加高斯噪声"""
    # 转换为YCbCr颜色空间
    ycbcr = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb).astype(np.float32) / 255.0
    
    # 添加噪声
    noise = np.random.normal(0, sigma, ycbcr.shape).astype(np.float32)
    ycbcr_noisy = np.clip(ycbcr + noise, 0, 1)
    
    # 转换回BGR颜色空间
    bgr_noisy = cv2.cvtColor((ycbcr_noisy * 255).astype(np.uint8), cv2.COLOR_YCrCb2BGR)
    return bgr_noisy

def add_high_frequency_noise(img, sigma):
    """添加高频噪声 - 更接近MATLAB的HighFN和ghp实现"""
    # 创建结果图像
    result = np.zeros_like(img, dtype=np.float32)
    
    # 对每个通道分别处理
    for i in range(3):
        # 计算傅里叶变换
        f = np.fft.fft2(img[:,:,i].astype(np.float32))
        
        # 创建高通滤波器 (类似于ghp.m)
        rows, cols = img.shape[:2]
        crow, ccol = rows // 2, cols // 2
        thresh = 10  # 与MATLAB中的thresh参数一致
        
        # 计算距离矩阵
        y, x = np.ogrid[:rows, :cols]
        d = np.sqrt((y - crow)**2 + (x - ccol)**2)
        
        # 创建高斯高通滤波器
        h = 1 - np.exp(-(d**2) / (2 * thresh**2))
        
        # 应用滤波器
        fshift = np.fft.fftshift(f)
        fshift = fshift * h
        f_ishift = np.fft.ifftshift(fshift)
        
        # 逆傅里叶变换
        img_back = np.fft.ifft2(f_ishift)
        img_back = np.abs(img_back)
        
        # 存储结果
        result[:,:,i] = img_back
    
    # 转换为uint8
    result = np.clip(result, 0, 255).astype(np.uint8)
    
    # 添加高斯噪声 (与MATLAB中的imnoise对应)
    noise = np.random.normal(0, sigma * 255, result.shape).astype(np.float32)
    result = np.clip(result.astype(np.float32) + noise, 0, 255).astype(np.uint8)
    
    return result

def add_impulse_noise(img, prob):
    """添加椒盐噪声"""
    output = np.copy(img)
    # 椒噪声
    black = np.random.random(img.shape[:2])
    output[black < prob/2] = 0
    # 盐噪声
    white = np.random.random(img.shape[:2])
    output[white < prob/2] = 255
    return output

def add_quantization_noise(img, levels):
    """添加量化噪声"""
    # 量化图像 - 更接近MATLAB的multithresh和imquantize实现
    result = np.zeros_like(img)
    
    # 对每个通道分别处理
    for i in range(3):
        # 计算直方图
        hist, bin_edges = np.histogram(img[:,:,i], bins=256, range=(0, 255))
        # 使用Otsu方法的扩展找到多个阈值
        thresholds = []
        # 简化版的多阈值Otsu
        for _ in range(levels-1):
            if len(thresholds) == 0:
                # 第一个阈值
                best_thresh = 0
                best_val = 0
                for t in range(1, 255):
                    w0 = np.sum(hist[:t]) / np.sum(hist)
                    w1 = 1 - w0
                    if w0 == 0 or w1 == 0:
                        continue
                    mu0 = np.sum(np.arange(t) * hist[:t]) / np.sum(hist[:t]) if np.sum(hist[:t]) > 0 else 0
                    mu1 = np.sum(np.arange(t, 256) * hist[t:]) / np.sum(hist[t:]) if np.sum(hist[t:]) > 0 else 0
                    val = w0 * w1 * (mu0 - mu1) ** 2
                    if val > best_val:
                        best_val = val
                        best_thresh = t
                thresholds.append(best_thresh)
            else:
                # 后续阈值 - 简化处理
                remaining = [t for t in range(1, 255) if all(abs(t-th) > 10 for th in thresholds)]
                if remaining:
                    thresholds.append(remaining[len(remaining)//2])
        
        # 排序阈值
        thresholds.sort()
        
        # 量化
        channel = img[:,:,i].copy()
        for j, thresh in enumerate(thresholds):
            mask = (channel <= thresh) & (channel > (0 if j == 0 else thresholds[j-1]))
            channel[mask] = 0 if j == 0 else thresholds[j-1]
        
        # 最后一个区间
        mask = channel > thresholds[-1]
        channel[mask] = 255
        
        result[:,:,i] = channel
    
    return result.astype(np.uint8)

def add_gaussian_blur(img, kernel_size):
    """添加高斯模糊"""
    # 确保kernel_size是奇数
    if kernel_size % 2 == 0:
        kernel_size += 1
    
    # 应用高斯模糊
    blurred = cv2.GaussianBlur(img, (kernel_size, kernel_size), kernel_size/6)
    return blurred

def apply_image_denoising(img, sigma):
    """使用BM3D进行图像去噪"""
    # 转换为浮点数并归一化
    img_float = img.astype(np.float32) / 255.0
    
    # 对每个通道应用BM3D
    if len(img.shape) == 3:
        result = np.zeros_like(img_float)
        for i in range(3):
            # 添加噪声
            noisy = img_float[:,:,i] + np.random.normal(0, sigma, img_float[:,:,i].shape)
            # 应用BM3D去噪
            result[:,:,i] = bm3d.bm3d(noisy, sigma)
    else:
        # 添加噪声
        noisy = img_float + np.random.normal(0, sigma, img_float.shape)
        # 应用BM3D去噪
        result = bm3d.bm3d(noisy, sigma)
    
    # 转换回uint8
    return np.clip(result * 255, 0, 255).astype(np.uint8)

def apply_jpeg_compression(img, output_path, quality):
    """应用JPEG压缩"""
    # 使用PIL保存为JPEG格式
    Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)).save(output_path, 'JPEG', quality=quality)

def apply_jp2k_compression(img, output_path, compression_ratio):
    """应用JPEG2000压缩"""
    # 使用OpenCV保存为JPEG2000格式
    cv2.imwrite(output_path, img, [cv2.IMWRITE_JPEG2000_COMPRESSION_X1000, int(compression_ratio * 1000)])

def add_non_eccentricity_pattern_noise(img, num_patches):
    """添加非偏心模式噪声"""
    result = img.copy()
    h, w = img.shape[:2]
    
    for _ in range(int(num_patches)):
        # 随机选择一个位置
        r_w = np.random.randint(0, w - 30)
        r_h = np.random.randint(0, h - 30)
        
        # 复制一个区域到另一个区域
        result[r_h:r_h+15, r_w:r_w+15] = img[r_h+5:r_h+20, r_w+5:r_w+20]
    
    return result

def apply_block_wise_distortion(img, block_size, level):
    """应用局部块状失真 - 更接近MATLAB的BlockWise实现"""
    # 计算各通道均值 
    r_mean = np.mean(img[:, :, 0]) 
    g_mean = np.mean(img[:, :, 1]) 
    b_mean = np.mean(img[:, :, 2]) 
    
    # 获取图像尺寸 
    h, w = img.shape[:2] 
    
    # 根据flag创建不同的block 
    if level == 5 or level == 4: 
        block = np.stack([np.ones((32, 32)) * r_mean, 
                         np.ones((32, 32)) * g_mean, 
                         np.ones((32, 32)) * b_mean], axis=2) 
    elif level == 3: 
        block = np.stack([np.ones((32, 32)) * (r_mean + 30), 
                         np.ones((32, 32)) * g_mean, 
                         np.ones((32, 32)) * b_mean], axis=2) 
    elif level == 2: 
        block = np.stack([np.ones((32, 32)) * (r_mean + 50), 
                         np.ones((32, 32)) * g_mean, 
                         np.ones((32, 32)) * b_mean], axis=2) 
    elif level == 1: 
        block = np.stack([np.zeros((32, 32)), 
                         np.ones((32, 32)) * g_mean, 
                         np.ones((32, 32)) * b_mean], axis=2) 
    
    # 转换为uint8类型 
    block = block.astype(np.uint8) 
    
    # 创建图像副本以避免修改原始图像 
    result_img = img.copy() 
    
    # 随机位置放置block 
    for i in range(block_size): 
        r_w = np.random.randint(0, h - 32)  # Python索引从0开始 
        r_h = np.random.randint(1, w - 32)  # 与MATLAB代码保持一致，h从1开始 
        result_img[r_w:r_w+32, r_h:r_h+32, :] = block 
    
    return result_img

def apply_mean_shift(img, shift):
    """应用均值偏移"""
    return np.clip(img.astype(np.float32) + shift, 0, 255).astype(np.uint8)

def apply_contrast_change(img, factor):
    """应用对比度变化"""
    return np.clip(img.astype(np.float32) * factor, 0, 255).astype(np.uint8)

def apply_color_saturation(img, factor):
    """应用色彩饱和度变化"""
    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
    
    # 调整饱和度
    hsv[:,:,1] = np.clip(hsv[:,:,1] * (1 + factor), 0, 255)
    
    # 转换回BGR颜色空间
    return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)

def add_multiplicative_gaussian_noise(img, sigma):
    """添加乘性高斯噪声"""
    result = img.copy().astype(np.float32)
    
    for i in range(3):
        noise = 1 + sigma * np.random.randn(img.shape[0], img.shape[1])
        result[:,:,i] = result[:,:,i] * noise
    
    return np.clip(result, 0, 255).astype(np.uint8)

def apply_color_quantization_dither(img, levels):
    """应用颜色量化抖动"""
    # 使用PIL进行颜色量化
    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    quantized = pil_img.quantize(colors=levels, dither=Image.FLOYDSTEINBERG)
    # 转换回BGR
    return cv2.cvtColor(np.array(quantized.convert('RGB')), cv2.COLOR_RGB2BGR)

def apply_color_aberrations(img, shift):
    """应用色差"""
    result = img.copy()
    
    # 分离通道
    b, g, r = cv2.split(result)
    
    # 对红色通道进行水平偏移
    rows, cols = r.shape
    M = np.float32([[1, 0, shift], [0, 1, 0]])
    r_shifted = cv2.warpAffine(r, M, (cols, rows))
    
    # 对蓝色通道进行水平偏移（一半的距离）
    M = np.float32([[1, 0, shift/2], [0, 1, 0]])
    b_shifted = cv2.warpAffine(b, M, (cols, rows))
    
    # 合并通道
    result = cv2.merge([b_shifted, g, r_shifted])
    
    # 应用轻微的高斯模糊
    result = cv2.GaussianBlur(result, (3, 3), 0.5)
    
    return result